import torch
from net import MyAlexNet
from torch.autograd import Variable
from torchvision import datasets, transforms
from torchvision.transforms import ToTensor
from torchvision.transforms import ToPILImage
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import random
import matplotlib.pyplot as plt

ROOT_TRAIN = r'data\train'
ROOT_TEST = r'data\val'

# 将图像的像素值归一化到【-1， 1】之间
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    normalize])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])

train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)

train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = MyAlexNet().to(device)

# 加载模型
model.load_state_dict(torch.load(r"G:\PY Project\AlexNet\save_model\best_model.pth"))

# 获取预测结果
classes = [
    "cat",
    "dog",
]

# 把张量转化为照片格式
show = ToPILImage()

# 进入到验证阶段
model.eval()

'''''
for i in a:
    x, y = val_dataset[i][0], val_dataset[i][1] #将一个样本的输入和标签分别存放在x和y两个变量中
    show(x).show()
    x = Variable(torch.unsqueeze(x, dim=0).float(), requires_grad=True).to(device)
    #为了将张量 x 转换为 PyTorch 中模型所需的输入形式，并在需要计算梯度的情况下，将其转换为可求梯度的 Variable 对象。
    x = torch.tensor(x).to(device)
    with torch.no_grad():
        pred = model(x)
        predicted, actual = classes[torch.argmax(pred[0])], classes[y]
        print(f'predicted:"{predicted}", Actual:"{actual}"')
'''''
# 定义一个用于显示图像和预测结果的函数
def show_prediction(x, y, model, classes, ax):
    # 将输入转换为 PyTorch Tensor 对象并移动到 GPU 上（如果有的话）
    x = x.to(device)
    # 关闭梯度计算，因为我们不会再训练过程中使用它
    with torch.no_grad():
        # 使用模型进行预测
        pred = model(x.unsqueeze(0))
        # 获取预测的类别和实际类别的索引
        predicted_cls_idx = torch.argmax(pred).item()
        actual_cls_idx = y
        # 获取预测类别和实际类别的文本标签
        predicted_cls = classes[predicted_cls_idx]
        actual_cls = classes[actual_cls_idx]

     # 反归一化
    x = x * 0.5 + 0.5
    ax.imshow(x.cpu().permute(1, 2, 0))
    ax.set_title(f'predicted:{predicted_cls}', fontsize=10)

# 随机选择一组索引
idx = random.sample(range(len(val_dataset)), 10)

fig, axes = plt.subplots(nrows=2, ncols=5, figsize=(10, 4))
fig.tight_layout()

# 遍历样本集，逐个显示图像和预测结果
for i, ax in enumerate(axes.flat):
    x, y = val_dataset[idx[i]][0], val_dataset[idx[i]][1]
    show_prediction(x, y, model, classes, ax)

plt.show()